from __future__ import division  
import numpy as np  
import scipy as sp  
from numpy.random import random
import pandas as pd

class SVD_CF:
    def __init__(self,X,k=20):
        
        self.X = X
        self.k = k
        self.mu = np.mean(self.X[:,2])
        
        self.bi = {}
        self.bu = {}
        self.qi = {}
        self.pu = {}
        
        self.ItemsForUser={}
        self.UsersForItem={}
        
        for i in range(self.X.shape[0]):  
            uid=self.X[i][0]
            i_id=self.X[i][1]
            rat=self.X[i][2]
            
            self.ItemsForUser.setdefault(i_id,{})  
            self.UsersForItem.setdefault(uid,{}) 
            
            self.ItemsForUser[i_id][uid]=rat  
            self.UsersForItem[uid][i_id]=rat  
            
            self.bi.setdefault(i_id,0)  
            self.bu.setdefault(uid,0)
            self.qi.setdefault(i_id,random((self.k,1))/10*(np.sqrt(self.k)))
            self.pu.setdefault(uid,random((self.k,1))/10*(np.sqrt(self.k)))
            
    def pred(self,uid,i_id):
        self.bi.setdefault(i_id,0)
        self.bu.setdefault(uid,0)
        
        self.qi.setdefault(i_id,np.zeros((self.k,1)))
        self.pu.setdefault(uid,np.zeros((self.k,1)))
        
        if(self.qi[i_id].all() == None):
            self.qi[i_id] = np.zeros((self.k,1))
        if(self.pu[uid].all() == None):
            self.pu[uid] = np.zeros((self.k,1))
            
        ans = self.mu + self.bi[i_id] + self.bu[uid] + np.sum(self.qi[i_id]*self.pu[uid])
        
        if ans > 5:
            return 5
        elif ans < 1:
            return 1
        return ans
    
    def train(self,steps = 25,gamma = 0.04,Lambda = 0.15):
        for step in range(steps):
            print("the %s-th step is running" %step)
            rmse_sum = 0.0
            
            kk = np.random.permutation(self.X.shape[0])
            for j in range(self.X.shape[0]):
                i = kk[j]
                uid = self.X[i][0]
                i_id = self.X[i][1]
                rat = self.X[i][2]
                
                eui = rat-self.pred(uid,i_id)
                
                rmse_sum += eui ** 2
                self.bu[uid]+=gamma*(eui-Lambda*self.bu[uid])  
                self.bi[i_id]+=gamma*(eui-Lambda*self.bi[i_id]) 
                
                temp=self.qi[i_id]  
                self.qi[i_id]+=gamma*(eui*self.pu[uid]-Lambda*self.qi[i_id])  
                self.pu[uid]+=gamma*(eui*temp-Lambda*self.pu[uid])
                
            gamma = gamma * 0.93
            print("the rmse of this step on train data is %s" %np.sqrt(rmse_sum/self.X.shape[0]))
            
    def test(self,test_X):
        output = pd.DataFrame(columns=['user', 'song', 'playcount', 'score', 'rank'])
        sums=0  
        test_X = np.array(test_X)  
          
        for i in range(test_X.shape[0]):
            pre=self.pred(test_X[i][0],test_X[i][1])
            output = output.append({'user': test_X[i][0],'song':test_X[i][1],'playcount':test_X[i][2],'score':pre}, ignore_index=True)
            sums+=(pre-test_X[i][2]) ** 2
        rmse=np.sqrt(sums/test_X.shape[0])
        print("the rmse on test data is %s" %rmse)  
        #进行排序，取前20
        output['rank'] = output['score'].rank(ascending=0, method='first')
        output = output.sort_values(by=['rank'])
        output = output.head(20)
        return output  
                